#include <iostream>
#include <fstream>
#include <vector>
#include "txtpara.h"
#include "KDindex.h"
#include "directsampling.h"
#include "tplate.h"
#include "taxonomy_tree.h"
#include "bayesian.h"

using namespace std;

inline float cube_root(float v)
{
    return pow(v,1.0/3);
}

int main(int argc, char *argv[])
{
    if (argc<6){
        cout<<"\nUsage:\n"<<argv[0]<<" param.cf ti.sgems model.txt template.par result\n\n";
        exit(0);
    }
    /* direct sampling parameters */
    int NR_neighbor;//Maximum number of neighbors
    int NR_try; //Maximum number of search
    float tolerance;//tolerance of difference
    /**/
    int xsize,ysize,zsize;
    double unitsize=1.0;
    char* config=argv[1];
    NR_neighbor=getparaINT("NR_neighbor",config);
    tolerance=getparaDOUBLE("tolerance",config);
    xsize=getparaINT("xsize",config);
    ysize=getparaINT("ysize",config);
    zsize=getparaINT("zsize",config);
    std::vector<coord> sample;
    //Step 1.  Generate sample points for the first time
    //>>>>>1.1 calculating Number of points
    std::vector <int> ids;
    double percentage = getparaFLOAT("percent", config);
    int tol_sample = xsize * ysize * zsize * percentage;
    int N_sample = xsize * ysize * zsize * percentage;
    int interval = round(sqrt(xsize * ysize * zsize / N_sample));
    if ((xsize/(interval))*(ysize/(interval))*(zsize/(interval)) > N_sample)
         interval++;
    //>>>>>1.2 generate random path
    srand(clock());
    int s_no = 0;
    for (int zpos = 0; zpos < zsize; zpos += interval)
    {
        for (int ypos = 0; ypos < ysize; ypos += interval)
        {
            for (int xpos = 0; xpos < xsize; xpos += interval)
            {
              coord s;
              s.x = xpos;
              s.y = ypos;
              s.z = zpos;
              sample.push_back(s);
              ids.push_back(s_no);
              s_no++;
            }
        }
    }
    N_sample = sample.size();  
    for (int zpos = 0; zpos < zsize; zpos++)
    {
        for (int ypos = 0; ypos < ysize; ypos++)
        {
            for (int xpos = 0; xpos < xsize; xpos++)
            {
              if(((xpos % interval)==0) && ((ypos % interval)==0) && ((zpos % interval)==0))
                  continue;
              coord s;              
              s.x = xpos;
              s.y = ypos;
              s.z = zpos;
              sample.push_back(s);
              ids.push_back(s_no);
              s_no++;
            }
        }
    }
    for(int i=N_sample;i<tol_sample;i++)
    {
        int ranid = N_sample+rand() % (xsize*ysize*zsize-N_sample);
        int tmp = ids[i];
        ids[i] = ids[ranid];
        ids[ranid] = tmp;
    }
    for (int s_no = 0; s_no < tol_sample; s_no++)
    {   
        int ranid = rand() % N_sample;
        int tmp = ids[s_no];
        ids[s_no] = ids[ranid];
        ids[ranid] = tmp;
    }
    N_sample=tol_sample;
    //Step 2.  The simulation on the coarsest grid
    Index <L2<float> > *R=createIndex();
    float *ti_ds;
    int tilength, tiwidth,tiheight;
    int start_id=0;
    int level=0;
    int *target=new int[int(xsize*ysize*zsize)];
    memset(target,-1,sizeof(int)*xsize*ysize*zsize);
    load_ti_3D(&ti_ds,argv[2],&tilength,&tiwidth,&tiheight);
    NR_try=int(getparaDOUBLE("Fraction",config)*tilength*tiwidth*tiheight);
    Simulation(R, ti_ds, NR_neighbor, NR_try, tolerance, sample, ids,
            target, N_sample, tilength, tiwidth, tiheight, unitsize, start_id, level,xsize,ysize,zsize);
    int maxscale, maxdv, NR_cate;
    int length, width, height;
    float *marginal_prob;
    float *marginal_prob_target;
    tplate tp;
    load_template_3D(&tp, argv[4]);
    height = zsize;
    length = xsize;
    width = ysize;
    maxdv = getparaINT("maxdv", config);
    NR_cate = getparaINT("NR_CATE", config);
    float *marginal = new float[NR_cate];
    maxscale = getparaINT("maxscale", config);
    marginal_prob = new float[NR_cate];
    marginal_prob_target = new float[NR_cate];
    getparaFLOATARRAY("marginal_prob", config, marginal_prob, NR_cate);
    getparaFLOATARRAY("target_marginal_prob", config, marginal_prob_target, NR_cate);
    float *logmarginal = new float[NR_cate];
    float *rectmarginal = new float[NR_cate];
    for (int c = 0; c < NR_cate; c++)
    {
        logmarginal[c] = log(marginal_prob[c]);
    }
    // Read the training image to the trianing grid
    ALL_taxonomy_tree *all_tree = new ALL_taxonomy_tree;
    load_taxonomy_tree(all_tree, tp.tpsize, argv[3]);
    train_grid_3D *ti;
    target_grid_3D *grid = new target_grid_3D;
    init_sim_grid_3D(grid, length, width, height, maxscale, NR_cate,
            N_sample, target);
    int *hard_cate = new int[NR_cate];
    int g=0;
    float *prob_result = new float[NR_cate];
    int cur_length = grid->length[g];
    int cur_width = grid->width[g];
    int cur_height = grid->height[g];

    for (int z = 0; z < cur_height; z++)
    {
        for (int y = 0; y < cur_width; y++)
        {
            for (int x = 0; x < cur_length; x++)
            {
                memcpy(hard_cate,
                        grid->sample[g].category + z*cur_width*cur_length*NR_cate+
                        y * cur_length * NR_cate + x * NR_cate,
                        NR_cate * sizeof(int));
                data_event *dv = new data_event;
                dv->tpos = new int[tp.tpsize];
                dv->pos = new int[tp.tpsize * 3];
                dv->value = new int *[tp.tpsize];
                get_dv_3D(dv, &tp, grid, g, x, y, z,maxdv);
                //if (0 == dv->dvsize)
                if (dv->dvsize <= 0)
                {  // no neighbors found, using the marginal pmf
                    memcpy(grid->sample[g].prob +
                            z * cur_width * cur_length *
                            NR_cate + y * cur_length * NR_cate +
                            x * NR_cate, marginal_prob,
                            sizeof(float) * NR_cate);
                }
                else
                {
                    marginal = marginal_prob_target;
                    memset(prob_result, 0,
                            sizeof(float) * NR_cate);
                    float *all_node_prob = extract_all_pdf(dv, all_tree->t_tree[g], marginal, NR_cate);
                    for (int c = 0; c < NR_cate; c++)
                    {
                        rectmarginal[c] = log(marginal[c]);
                    }
                    for (int ii = 0; ii < dv->dvsize; ii++)
                    {
                        float sum = 0;
                        for (int c = 0; c < NR_cate; c++)
                        {
                            sum += all_node_prob[ii * 2 + c];
                        }
                        for (int c = 0; c < NR_cate; c++)
                        {
                            all_node_prob[ii * 2 + c] /= sum;
                        }
                    }
                    merge_prob(prob_result, all_node_prob, logmarginal, rectmarginal, NR_cate, dv->dvsize); 
                    memcpy(grid->sample[g].prob + z* cur_width*cur_length*NR_cate + y * cur_length * NR_cate + x * NR_cate, prob_result, sizeof(float) * NR_cate);
                }
                for (int i = 0; i < dv->dvsize; i++)
                {
                    delete dv->value[i];
                }
                delete dv->value;
                delete dv->tpos;
                delete dv->pos;
                delete dv;
            }
        }
    }
    
    char tmpstr[256];
    sprintf(tmpstr, "%s.sgems", argv[5]);
    ofstream ofs(tmpstr);
    ofs<<int(xsize)<<" "<<int(ysize)<<" "<<int(zsize)<<endl;
    ofs<<1<<endl;
    ofs<<"facies"<<endl;
    for (int z = 0; z < zsize; z++)
    {
        for (int y = 0; y < ysize; y++)
        {
            for (int x = 0; x < xsize; x++)
            {
                float maxp=0;
                int maxc=0;
                for (int c=0;c<NR_cate;c++){
                    float cur_prob=grid->sample[g].prob[z * grid->width[g] * grid->length[g] * NR_cate + y * grid->length[g] * NR_cate + x * NR_cate + c];
                    if(cur_prob>=maxp){
                        maxp=cur_prob;
                        maxc=c;
                    }
                }
                ofs<<maxc<<" ";
            }
            ofs << endl;
        }
        ofs << endl;
    }
    delete[]hard_cate;
    return 0;
}
